# $Id: rcm_sim_split_pop1.r,v 1.1 2006/03/17 00:30:51 jkatz Exp $
# rcm_sim_split_pop1.r
# by Jonathan N. Katz <jkatz@caltech.edu>
#
#
#############################################################
##                                                         ##
##      Load Libraries                                     ##
##                                                         ##
#############################################################
library(nlme)

#############################################################
##                                                         ##
##      User Setable Parameters                            ##
##                                                         ##
#############################################################
iter       <- 1000          # Number of simulation iterations 
n.groups   <- 20            # Number of units
n.reps     <- 20            # Number observations per group
b.mean     <- 5             # Mean of \beta_i
b.std      <- 0             # Standard deviation of \beta_i           
dif.start  <- 0             # Starting difference in first two units
dif.end    <- 5             # Ending differecne in first two units
dif.step   <- 0.5           # Step size for  difference in first two units
set.seed(666)          # Seed for RNG

#############################################################
##                                                         ##
##      Generate objects not changed during Design loop    ##
##                                                         ##
#############################################################

## Storage Arrays
design.length <-  length(seq(dif.start,dif.end,by=dif.step))

rmse.colnames  <-  c("b.dif", "b.mean","b.std", "b.group")

rmse.lm               <- array(NA, c(design.length,4))
colnames(rmse.lm)     <- rmse.colnames

rmse.lmList           <- array(NA, c(design.length,4))
colnames(rmse.lmList) <-  rmse.colnames

rmse.lmeML            <- array(NA, c(design.length,4))
colnames(rmse.lmeML)  <-  rmse.colnames

rmse.lmeREML          <- array(NA, c(design.length,4))
colnames(rmse.lmeREML)<-  rmse.colnames


############################################################
##                                                         ##
##      Design Loop                                        ##
##                                                         ##
#############################################################

j <-  1
for (b.dif in seq(dif.start,dif.end,by=dif.step)) {
  cat("\n", "Expermiment, difference in first two units ", b.dif,"\n" )

#############################################################
##                                                         ##
##      Generate objects not changed during MC looping     ##
##                                                         ##
#############################################################
  
### Useful fixed values
  n         <- n.groups*n.reps
  group.id  <- rep (1:n.groups, each=n.reps)
 
### Generate \beta_i's for all runs
  b.group.all <-  array(rnorm(n.groups*iter, mean=b.mean, sd=b.std),
                      c(iter,n.groups))
  b.group.all[,1:2] <- b.group.all[,1:2] + b.dif

#### Arrays to store results
  
  results.colnames          <-  c("b.mean","b.std",paste("b",1:n.groups,sep=""))
  
  results.lm                <- array(NA,c(iter,n.groups+2))
  colnames(results.lm)      <- results.colnames
  
  results.lmList            <- array(NA,c(iter,n.groups+2))  
  colnames(results.lmList)  <- results.colnames
  
 
  results.lmeML             <- array(NA,c(iter,n.groups+2))
  colnames(results.lmeML)   <- results.colnames
  
  results.lmeREML           <- array(NA,c(iter,n.groups+2))
  colnames(results.lmeREML) <- results.colnames
  
#############################################################
##                                                         ##
##      Monte Carlo Loop                                   ##
##                                                         ##
#############################################################

  for (i in 1:iter) {
    
### Where are we in the loop
    cat(i, " ")    

    
### Generate X's
    x <- rnorm(n,mean=1,sd=0.10)

### Generate Y's
    b.group  <- b.group.all[i,]
    y        <- b.group[group.id]*x + rnorm(n, mean=0, sd=1)
    sim.data <- data.frame(y=y,x=x,group.id=as.factor(group.id))
    
### Pooled OLS
    sim.lm <-  lm(y ~  x, data=sim.data)
    results.lm[i,] <- unlist(c(coef(sim.lm)[2],NA,
                               rep(coef(sim.lm)[2],n.groups)))
    
### Unit by Unit OLS              
    sim.lmList <-  lm(y ~ group.id/(x-1), data=sim.data)
    results.lmList[i,] <-  unlist(c(NA,NA,
                                    coef(sim.lmList)[(n.groups+1):(2*n.groups)]
                                    ))

### RCM via ML
    sim.lmeML <-  lme(y ~ x, random= ~ 0 + x  | group.id, data=sim.data,
                      method="ML")
    results.lmeML[i,] <-  unlist(c(fixef(sim.lmeML)[2],
                                 sqrt(as.numeric(getVarCov(sim.lmeML))),
                                   coef(sim.lmeML)[,2]))
    
### RCM via REML
    sim.lmeREML <-  lme(y ~ x ,random= ~ 0 + x | group.id, data=sim.data,
                        method="REML")
    results.lmeREML[i,] <-  unlist(c(fixef(sim.lmeREML)[2],
                                     sqrt(as.numeric(getVarCov(sim.lmeREML))),
                                     coef(sim.lmeREML)[,2]))
  }
  
  
#### Calculate RMSE
  
# This cose creates and expands the true parameters to a matrix the
# same size as the results.

  truth <-  c(b.mean,b.std)
  truth <-  t(array(truth,c(2,iter)))
  truth <-  cbind(truth,b.group.all)
  
  mse.lm           <- apply((results.lm-truth)^2,2,mean)
  rmse.lm[j,]      <- unlist(c(b.dif, sqrt(mse.lm[1:2]),
                               sqrt(mean(mse.lm[3:length(mse.lm)]))))
  
  mse.lmList       <- apply((results.lmList-truth)^2,2,mean)
  rmse.lmList[j,]  <- unlist(c(b.dif, sqrt(mse.lmList[1:2]),
                               sqrt(mean(mse.lmList[3:length(mse.lmList)]))))
  
  mse.lmeML        <- apply((results.lmeML-truth)^2,2,mean)
  rmse.lmeML[j,]   <- unlist(c(b.dif, sqrt(mse.lmeML[1:2]),
                               sqrt(mean(mse.lmeML[3:length(mse.lmeML)]))))
  
  mse.lmeREML      <- apply((results.lmeREML-truth)^2,2,mean)
  rmse.lmeREML[j,] <- unlist(c(b.dif, sqrt(mse.lmeREML[1:2]),
                               sqrt(mean(mse.lmeREML[3:length(mse.lmeREML)]))))

  save(rmse.lm, rmse.lmList, rmse.lmeML, rmse.lmeREML,
       file="rmse_sim_split_pop1.Rdata", ascii=TRUE)

j <-  j+1
}


# In order to do the FGLS estimate, extract the vcov matrix using
# lapply(test.lis, vcov)

# Trick to get the se from the lme() 
# summary(sim.lmeML)$tTable[2]
 #do something like save(rmse.lm, rmse.,,, file="rsme.sim.Rdata")
